import json
import os
import cv2
from pathlib import Path

def convert_to_yolo(size, box):
    """将边界框坐标转换为 YOLO 格式"""
    dw = 1. / size[0]  # 宽度归一化
    dh = 1. / size[1]  # 高度归一化
    x = (box[0] + box[2]) / 2.0  # 中心 x
    y = (box[1] + box[3]) / 2.0  # 中心 y
    w = box[2] - box[0]  # 宽度
    h = box[3] - box[1]  # 高度
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def get_all_classes_from_json_files(json_files):
    """从所有JSON文件中自动提取类别列表"""
    classes = set()
    for json_file in json_files:
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            for shape in data.get('shapes', []):
                classes.add(shape.get('label', ''))
        except Exception as e:
            print(f"警告: 读取 {json_file} 时出错: {e}")
    return sorted(list(classes))

def find_image_file(image_filename, labels_dir):
    """在对应的images目录中查找图片文件"""
    # 从labels目录推导出images目录
    # 例如: F:\RDK_Tool_Train_Dataset\labels\train -> F:\RDK_Tool_Train_Dataset\images\train
    path_parts = labels_dir.replace('\\', '/').split('/')
    
    # 找到labels部分并替换为images
    if 'labels' in path_parts:
        labels_index = path_parts.index('labels')
        path_parts[labels_index] = 'images'
        images_dir = '/'.join(path_parts).replace('/', os.sep)
    else:
        # 如果没有找到labels，假设在同一目录
        images_dir = labels_dir
    
    print(f"查找图片目录: {images_dir}")
    
    # 去掉扩展名，只保留文件名
    base_name = os.path.splitext(os.path.basename(image_filename))[0]
    
    # 支持的图片扩展名
    extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
    
    # 在images目录中查找
    for ext in extensions:
        potential_path = os.path.join(images_dir, base_name + ext)
        if os.path.exists(potential_path):
            print(f"找到图片: {potential_path}")
            return potential_path
    
    # 如果没找到，尝试原始文件名
    original_filename = os.path.basename(image_filename)
    original_path = os.path.join(images_dir, original_filename)
    if os.path.exists(original_path):
        print(f"找到图片: {original_path}")
        return original_path
    
    print(f"未找到图片文件: {base_name}")
    return None

def json_to_yolo_converter(json_file, classes, labels_dir):
    """处理单个JSON文件并生成TXT文件"""
    try:
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
          # 获取图像路径
        image_filename = data['imagePath']
        
        # 如果imagePath是相对路径，需要解析实际路径
        if '..' in image_filename or '\\' in image_filename or '/' in image_filename:
            # 解析相对路径
            json_dir = os.path.dirname(json_file)
            image_full_path = os.path.join(json_dir, image_filename)
            image_full_path = os.path.normpath(image_full_path)
            
            # 如果解析的路径存在，使用它
            if os.path.exists(image_full_path):
                image_path = image_full_path
                print(f"使用相对路径找到图片: {image_path}")
            else:
                # 否则尝试在对应的images目录中查找
                image_basename = os.path.basename(image_filename)
                image_path = find_image_file(image_basename, labels_dir)
        else:
            # 如果是简单文件名，在对应的images目录中查找
            image_path = find_image_file(image_filename, labels_dir)
        
        if image_path is None:
            print(f"警告: 找不到对应的图像文件 {image_filename}")
            return False
        
        print(f"正在处理: {os.path.basename(json_file)}")
        
        # 加载图片获取尺寸
        image = cv2.imread(image_path)
        if image is None:
            print(f"警告: 无法加载图像 {image_path}")
            return False
        
        height, width = image.shape[:2]
        
        # 生成输出文件路径（与JSON文件在同一目录）
        json_basename = os.path.basename(json_file)
        txt_filename = json_basename.replace('.json', '.txt')
        out_file = os.path.join(labels_dir, txt_filename)
        
        # 检查是否已经存在对应的txt文件
        if os.path.exists(out_file):
            print(f"跳过: {txt_filename} 已存在")
            return True
        
        # 转换标注并写入文件
        with open(out_file, 'w') as f:
            for shape in data['shapes']:
                label = shape['label']
                if label not in classes:
                    print(f"跳过未知类别: {label}")
                    continue
                
                cls_id = classes.index(label)
                points = shape['points']
                
                # 处理矩形标注
                if len(points) == 2:  # 矩形标注
                    xmin = min(points[0][0], points[1][0])
                    xmax = max(points[0][0], points[1][0])
                    ymin = min(points[0][1], points[1][1])
                    ymax = max(points[0][1], points[1][1])
                else:  # 多边形标注，取外接矩形
                    x_coords = [point[0] for point in points]
                    y_coords = [point[1] for point in points]
                    xmin = min(x_coords)
                    xmax = max(x_coords)
                    ymin = min(y_coords)
                    ymax = max(y_coords)
                
                box = (xmin, ymin, xmax, ymax)
                bb = convert_to_yolo((width, height), box)
                f.write(f"{cls_id} {' '.join([str(a) for a in bb])}\n")
        
        print(f"转换完成: {txt_filename}")
        return True
        
    except Exception as e:
        print(f"转换失败 {json_file}: {str(e)}")
        return False

def convert_folder_json_to_yolo(folder_path):
    """
    将指定文件夹中的所有LabelMe JSON文件转换为YOLO格式
    
    Args:
        folder_path (str): 包含JSON文件和图片的文件夹路径
    """
    print("=" * 60)
    print("LabelMe JSON 转 YOLO TXT 格式工具")
    print("=" * 60)
    print(f"目标文件夹: {folder_path}")
    print("-" * 60)
    
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"❌ 错误: 文件夹不存在 {folder_path}")
        return
    
    # 获取所有JSON文件
    json_files = []
    for file in os.listdir(folder_path):
        if file.endswith('.json'):
            json_files.append(os.path.join(folder_path, file))
    
    if not json_files:
        print(f"❌ 在 {folder_path} 中没有找到JSON文件")
        return
    
    print(f"📂 找到 {len(json_files)} 个JSON文件")
    
    # 自动提取类别列表
    print("🔍 正在自动提取类别...")
    classes = get_all_classes_from_json_files(json_files)
    
    if not classes:
        print("❌ 没有找到任何类别标签")
        return
    
    print(f"📋 发现类别: {classes}")
    print("-" * 60)
    
    # 检查有多少文件已经转换过
    existing_txt_count = 0
    for json_file in json_files:
        json_basename = os.path.basename(json_file)
        txt_filename = json_basename.replace('.json', '.txt')
        txt_path = os.path.join(folder_path, txt_filename)
        if os.path.exists(txt_path):
            existing_txt_count += 1
    
    if existing_txt_count > 0:
        print(f"ℹ️  已存在 {existing_txt_count} 个TXT文件，将跳过这些文件")
        print("-" * 60)
    
    # 批量转换所有JSON文件
    success_count = 0
    skipped_count = 0
    
    for json_file in json_files:
        json_basename = os.path.basename(json_file)
        txt_filename = json_basename.replace('.json', '.txt')
        txt_path = os.path.join(folder_path, txt_filename)
        
        # 如果txt文件已存在，跳过
        if os.path.exists(txt_path):
            print(f"⏭️  跳过: {txt_filename} (已存在)")
            skipped_count += 1
            continue
        
        if json_to_yolo_converter(json_file, classes, folder_path):
            success_count += 1
    
    print("-" * 60)
    print(f"✅ 转换完成!")
    print(f"📊 成功转换: {success_count} 个文件")
    print(f"⏭️  跳过文件: {skipped_count} 个文件")
    print(f"📁 总文件数: {len(json_files)} 个文件")
    
    # 生成classes.txt文件
    classes_file = os.path.join(folder_path, 'classes.txt')
    with open(classes_file, 'w', encoding='utf-8') as f:
        for cls in classes:
            f.write(f"{cls}\n")
    print(f"📋 类别文件已生成: {classes_file}")
    print("=" * 60)

def main():
    """主函数 - 简化的用户交互"""
    print("LabelMe JSON 转 YOLO TXT 格式工具")
    print("=" * 40)
    
    # 默认路径
    default_path = r'F:\RDK_Tool_Train_Dataset\labels\val'
    
    print(f"默认路径: {default_path}")
    user_input = input("请输入文件夹路径 (直接回车使用默认路径): ").strip()
    
    # 如果用户没有输入，使用默认路径
    if not user_input:
        folder_path = default_path
    else:
        folder_path = user_input.strip('"\'')  # 去掉可能的引号
    
    # 转换路径格式
    folder_path = os.path.abspath(folder_path)
    
    print(f"\n将处理文件夹: {folder_path}")
    confirm = input("确认开始转换? (y/n，直接回车默认为y): ").strip().lower()
    
    if confirm == 'n':
        print("已取消操作")
        return
    
    # 执行转换
    convert_folder_json_to_yolo(folder_path)

if __name__ == "__main__":
    main()